/** 
 *  Copyright 2010 Matthew Werny & Burk Price.
 * 
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 * 
 *       http://www.apache.org/licenses/LICENSE-2.0
 * 
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *  under the License.
 */
package net.iceflow.db.migration.manager;

import net.iceflow.db.migration.resolver.ResourceMigrationResolver;
import net.iceflow.db.migration.version.strategy.SimpleVersionStrategy;
import net.iceflow.db.migration.version.VersionStrategy;
import net.iceflow.db.migration.exception.MigrationException;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Date;
import java.util.List;
import java.util.Set;
import java.util.SortedSet;
import java.util.TreeSet;

import javax.sql.DataSource;

import net.iceflow.db.jdbc.DatabaseType;
import net.iceflow.db.jdbc.DatabaseUtils;
import net.iceflow.db.migration.MigrationExecuter;
import net.iceflow.db.migration.MigrationManager;
import net.iceflow.db.migration.MigrationResolver;

import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.collections.Predicate;
import org.apache.commons.lang.time.DurationFormatUtils;
import org.apache.commons.lang.time.StopWatch;
import org.apache.log4j.Logger;
import org.springframework.dao.DataAccessException;
import org.springframework.jdbc.core.ConnectionCallback;
import org.springframework.jdbc.core.JdbcTemplate;

public class DataSourceMigrationManager implements MigrationManager {

    protected final Logger logger = Logger.getLogger(getClass());
    private final JdbcTemplate jdbcTemplate;
    private DatabaseType dbType;
    private VersionStrategy versionStrategy = new SimpleVersionStrategy();
    private MigrationResolver migrationResolver = new ResourceMigrationResolver();

    public DataSourceMigrationManager(DataSource dataSource) {
        this.jdbcTemplate = new JdbcTemplate(dataSource);
        this.dbType = determineDatabaseType();
    }

    public DataSourceMigrationManager(DataSource dataSource, DatabaseType dbType) {
        this(dataSource);
        this.dbType = dbType;
    }

    protected void enableMigrations() {
        try {
            jdbcTemplate.execute(new ConnectionCallback() {

                public Object doInConnection(Connection connection) throws SQLException, DataAccessException {
                    versionStrategy.enableVersioning(dbType, connection);
                    return null;
                }
            });

            logger.info("Successfully enabled migrations.");
        } catch (DataAccessException e) {
            logger.error("Could not enable migrations.", e);
            throw new MigrationException(e);
        }
    }

    public boolean validate() {
        return pendingMigrations().isEmpty();
    }

    public SortedSet<MigrationExecuter> pendingMigrations() {
        Set<String> appliedMigrations = determineAppliedMigrationVersions();
        Set<MigrationExecuter> availableMigrations = migrationResolver.resolve(dbType);

        SortedSet<MigrationExecuter> pendingMigrations = new TreeSet<MigrationExecuter>();
        CollectionUtils.select(availableMigrations,
                new PendingMigrationPredicate(appliedMigrations),
                pendingMigrations);

        return pendingMigrations;
    }

    public void migrate() {
        Set<String> appliedMigrations = determineAppliedMigrationVersions();

        if (appliedMigrations == null) {
            enableMigrations();
            appliedMigrations = Collections.EMPTY_SET;
        }

        Set<MigrationExecuter> availableMigrations = migrationResolver.resolve(dbType);

        // Which migrations need to be applied (ie: are pending)?
        // TODO This seems like something that could be in its own method.
        final List<MigrationExecuter> pendingMigrations = new ArrayList<MigrationExecuter>(availableMigrations.size());
        CollectionUtils.select(availableMigrations, new PendingMigrationPredicate(appliedMigrations), pendingMigrations);
        Collections.sort(pendingMigrations);

        if (pendingMigrations.isEmpty()) {
            logger.info("Database is up to date; no migration necessary.");
            return;
        }

        // TODO Check that no two pending migrations are the same version.
        StopWatch watch = new StopWatch();
        watch.start();

        logger.info("Migrating database... applying "
                + pendingMigrations.size() + " migration"
                + (pendingMigrations.size() > 1 ? "s" : "") + ".");

        try {
            jdbcTemplate.execute(new ConnectionCallback() {

                public Object doInConnection(Connection connection)
                        throws SQLException, DataAccessException {
                    int successfulCount = 0;
                    MigrationExecuter currentMigration = null;

                    final boolean autoCommit = connection.getAutoCommit();
                    connection.setAutoCommit(false);

                    try {
                        for (MigrationExecuter migration : pendingMigrations) {
                            currentMigration = migration;
                            logger.info("Running migration "
                                    + currentMigration.getFilename() + ".");

                            final Date startTime = new Date();
                            StopWatch migrationWatch = new StopWatch();
                            migrationWatch.start();

                            currentMigration.migrate(dbType, connection);
                            versionStrategy.recordMigration(dbType, connection,
                                    currentMigration.getVersion(), startTime,
                                    migrationWatch.getTime());

                            connection.commit();

                            ++successfulCount;
                        }
                    } catch (Throwable e) {
                        assert currentMigration != null;
                        String message = "Migration for version "
                                + currentMigration.getVersion()
                                + " failed, rolling back and terminating migration.";
                        logger.error(message, e);
                        connection.rollback();
                        throw new MigrationException(message, e);
                    } finally {
                        connection.setAutoCommit(autoCommit);
                    }

                    return successfulCount;
                }
            });
        } catch (DataAccessException e) {
            logger.error("Failed to migrate database.", e);
            throw new MigrationException(e);
        }

        watch.stop();

        logger.info("Migrated database in "
                + DurationFormatUtils.formatDurationHMS(watch.getTime()) + ".");
    }

    public void setDatabaseType(DatabaseType dbType) {
        this.dbType = dbType;
    }

    public void setMigrationResolver(MigrationResolver migrationResolver) {
        this.migrationResolver = migrationResolver;
    }

    public void setVersionStrategy(VersionStrategy versionStrategy) {
        this.versionStrategy = versionStrategy;
    }

    private DatabaseType determineDatabaseType() {
        return (DatabaseType) jdbcTemplate.execute(new ConnectionCallback() {

            public Object doInConnection(Connection connection) throws SQLException, DataAccessException {
                return DatabaseUtils.databaseType(connection.getMetaData().getURL());
            }
        });
    }

    private Set<String> determineAppliedMigrationVersions() {
        return (Set<String>) jdbcTemplate.execute(new ConnectionCallback() {

            public Object doInConnection(Connection connection) throws SQLException, DataAccessException {
                return versionStrategy.appliedMigrations(dbType, connection);
            }
        });
    }

    private static class PendingMigrationPredicate implements Predicate {

        private final Set<String> appliedMigrations;

        public PendingMigrationPredicate(Set<String> appliedMigrations) {
            this.appliedMigrations = appliedMigrations == null ? Collections.EMPTY_SET : appliedMigrations;
        }

        public boolean evaluate(Object input) {
            if (input instanceof MigrationExecuter) {
                return !appliedMigrations.contains(((MigrationExecuter) input).getVersion());
            } else {
                return !appliedMigrations.contains(input.toString());
            }
        }
    }
}
